{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d792fb5b",
   "metadata": {},
   "source": [
    "<picture>\n",
    "  <source media=\"(prefers-color-scheme: dark)\" srcset=\"https://vespa.ai/assets/vespa-ai-logo-heather.svg\">\n",
    "  <source media=\"(prefers-color-scheme: light)\" srcset=\"https://vespa.ai/assets/vespa-ai-logo-rock.svg\">\n",
    "  <img alt=\"#Vespa\" width=\"200\" src=\"https://vespa.ai/assets/vespa-ai-logo-rock.svg\" style=\"margin-bottom: 25px;\">\n",
    "</picture>\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "23cf319b",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/vector_stores/VespaIndexDemo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "307804a3-c02b-4a57-ac0d-172c30ddc851",
   "metadata": {},
   "source": [
    "# Vespa Vector Store demo\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5508d8ac",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0beb6603",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-vector-stores-vespa llama-index pyvespa"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f7010b1d-d1bb-4f08-9309-a328bb4ea396",
   "metadata": {},
   "source": [
    "#### Setting up API key\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08ad68ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import openai\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8ee4473a-094f-4d0a-a825-e1213db07240",
   "metadata": {},
   "source": [
    "#### Load documents, build the VectorStoreIndex\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2bcc07",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.vector_stores.vespa import VespaVectorStore\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3a41a70d",
   "metadata": {},
   "source": [
    "## Defining some sample data\n",
    "\n",
    "Let's insert some documents.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df6b6d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.schema import TextNode\n",
    "\n",
    "nodes = [\n",
    "    TextNode(\n",
    "        text=\"The Shawshank Redemption\",\n",
    "        metadata={\n",
    "            \"author\": \"Stephen King\",\n",
    "            \"theme\": \"Friendship\",\n",
    "            \"year\": 1994,\n",
    "        },\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=\"The Godfather\",\n",
    "        metadata={\n",
    "            \"director\": \"Francis Ford Coppola\",\n",
    "            \"theme\": \"Mafia\",\n",
    "            \"year\": 1972,\n",
    "        },\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=\"Inception\",\n",
    "        metadata={\n",
    "            \"director\": \"Christopher Nolan\",\n",
    "            \"theme\": \"Fiction\",\n",
    "            \"year\": 2010,\n",
    "        },\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=\"To Kill a Mockingbird\",\n",
    "        metadata={\n",
    "            \"author\": \"Harper Lee\",\n",
    "            \"theme\": \"Mafia\",\n",
    "            \"year\": 1960,\n",
    "        },\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=\"1984\",\n",
    "        metadata={\n",
    "            \"author\": \"George Orwell\",\n",
    "            \"theme\": \"Totalitarianism\",\n",
    "            \"year\": 1949,\n",
    "        },\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=\"The Great Gatsby\",\n",
    "        metadata={\n",
    "            \"author\": \"F. Scott Fitzgerald\",\n",
    "            \"theme\": \"The American Dream\",\n",
    "            \"year\": 1925,\n",
    "        },\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=\"Harry Potter and the Sorcerer's Stone\",\n",
    "        metadata={\n",
    "            \"author\": \"J.K. Rowling\",\n",
    "            \"theme\": \"Fiction\",\n",
    "            \"year\": 1997,\n",
    "        },\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31fe4378",
   "metadata": {},
   "source": [
    "### Initilizing the VespaVectorStore\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0be7d09",
   "metadata": {},
   "source": [
    "To make it really simple to get started, we provide a template Vespa application that will be deployed upon initializing the vector store.\n",
    "\n",
    "This is a huge abstraction and there are endless opportunities to tailor and customize the Vespa application to your needs. But for now, let's keep it simple and initialize with the default template.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30b0b2e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import StorageContext\n",
    "\n",
    "vector_store = VespaVectorStore()\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex(nodes, storage_context=storage_context)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71a4a3ec",
   "metadata": {},
   "source": [
    "### Deleting documents\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4637a79",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_to_delete = nodes[0].node_id\n",
    "node_to_delete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84a97903",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store.delete(ref_doc_id=node_to_delete)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03315550",
   "metadata": {},
   "source": [
    "## Querying\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74cabf95",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.vector_stores.types import (\n",
    "    VectorStoreQuery,\n",
    "    VectorStoreQueryMode,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6401e25",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = VectorStoreQuery(\n",
    "    query_str=\"Great Gatsby\",\n",
    "    mode=VectorStoreQueryMode.TEXT_SEARCH,\n",
    "    similarity_top_k=1,\n",
    ")\n",
    "result = vector_store.query(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09f1bf81",
   "metadata": {},
   "outputs": [],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77d2528e",
   "metadata": {},
   "source": [
    "## As retriever\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8d7aca1",
   "metadata": {},
   "source": [
    "### Default query mode (text search)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a71818e",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = index.as_retriever(vector_store_query_mode=\"default\")\n",
    "results = retriever.retrieve(\"Who directed inception?\")\n",
    "display(Markdown(f\"**Retrieved nodes:**\\n {results}\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfe83ebf",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = index.as_retriever(vector_store_query_mode=\"semantic_hybrid\")\n",
    "results = retriever.retrieve(\"Who wrote Harry Potter?\")\n",
    "display(Markdown(f\"**Retrieved nodes:**\\n {results}\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8aa36e8",
   "metadata": {},
   "source": [
    "### As query engine\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1bd18f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\"Who directed inception?\")\n",
    "display(Markdown(f\"**Response:** {response}\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aede9cf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    vector_store_query_mode=\"semantic_hybrid\", verbose=True\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"When was the book about the wizard boy published and what was it called?\"\n",
    ")\n",
    "display(Markdown(f\"**Response:** {response}\"))\n",
    "display(Markdown(f\"**Sources:** {response.source_nodes}\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90081efd",
   "metadata": {},
   "source": [
    "## Using metadata filters\n",
    "\n",
    "**NOTE**: This metadata filtering is done by llama-index, outside of vespa. For native and much more performant filtering, you should use Vespa's own filtering capabilities.\n",
    "\n",
    "See [Vespa's documentation](https://docs.vespa.ai/en/reference/query-language-reference.html) for more information.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0663ab38",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.vector_stores import (\n",
    "    FilterOperator,\n",
    "    FilterCondition,\n",
    "    MetadataFilter,\n",
    "    MetadataFilters,\n",
    ")\n",
    "\n",
    "# Let's define a filter that will only allow nodes that has the theme \"Fiction\" OR is published after 1997\n",
    "\n",
    "filters = MetadataFilters(\n",
    "    filters=[\n",
    "        MetadataFilter(key=\"theme\", value=\"Fiction\"),\n",
    "        MetadataFilter(key=\"year\", value=1997, operator=FilterOperator.GT),\n",
    "    ],\n",
    "    condition=FilterCondition.OR,\n",
    ")\n",
    "\n",
    "retriever = index.as_retriever(filters=filters)\n",
    "result = retriever.retrieve(\"Harry Potter\")\n",
    "display(Markdown(f\"**Result:** {result}\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "414e6d78",
   "metadata": {},
   "source": [
    "## Abstraction level of this integration\n",
    "\n",
    "To make it really simple to get started, we provide a template Vespa application that will be deployed upon initializing the vector store. This removes some of the complexity of setting up Vespa for the first time, but for serious use cases, we strongly recommend that you read the [Vespa documentation](docs.vespa.ai) and tailor the application to your needs.\n",
    "\n",
    "### The template\n",
    "\n",
    "The provided template Vespa application can be seen below:\n",
    "\n",
    "```python\n",
    "from vespa.package import (\n",
    "    ApplicationPackage,\n",
    "    Field,\n",
    "    Schema,\n",
    "    Document,\n",
    "    HNSW,\n",
    "    RankProfile,\n",
    "    Component,\n",
    "    Parameter,\n",
    "    FieldSet,\n",
    "    GlobalPhaseRanking,\n",
    "    Function,\n",
    ")\n",
    "\n",
    "hybrid_template = ApplicationPackage(\n",
    "    name=\"hybridsearch\",\n",
    "    schema=[\n",
    "        Schema(\n",
    "            name=\"doc\",\n",
    "            document=Document(\n",
    "                fields=[\n",
    "                    Field(name=\"id\", type=\"string\", indexing=[\"summary\"]),\n",
    "                    Field(name=\"metadata\", type=\"string\", indexing=[\"summary\"]),\n",
    "                    Field(\n",
    "                        name=\"text\",\n",
    "                        type=\"string\",\n",
    "                        indexing=[\"index\", \"summary\"],\n",
    "                        index=\"enable-bm25\",\n",
    "                        bolding=True,\n",
    "                    ),\n",
    "                    Field(\n",
    "                        name=\"embedding\",\n",
    "                        type=\"tensor<float>(x[384])\",\n",
    "                        indexing=[\n",
    "                            \"input text\",\n",
    "                            \"embed\",\n",
    "                            \"index\",\n",
    "                            \"attribute\",\n",
    "                        ],\n",
    "                        ann=HNSW(distance_metric=\"angular\"),\n",
    "                        is_document_field=False,\n",
    "                    ),\n",
    "                ]\n",
    "            ),\n",
    "            fieldsets=[FieldSet(name=\"default\", fields=[\"text\", \"metadata\"])],\n",
    "            rank_profiles=[\n",
    "                RankProfile(\n",
    "                    name=\"bm25\",\n",
    "                    inputs=[(\"query(q)\", \"tensor<float>(x[384])\")],\n",
    "                    functions=[Function(name=\"bm25sum\", expression=\"bm25(text)\")],\n",
    "                    first_phase=\"bm25sum\",\n",
    "                ),\n",
    "                RankProfile(\n",
    "                    name=\"semantic\",\n",
    "                    inputs=[(\"query(q)\", \"tensor<float>(x[384])\")],\n",
    "                    first_phase=\"closeness(field, embedding)\",\n",
    "                ),\n",
    "                RankProfile(\n",
    "                    name=\"fusion\",\n",
    "                    inherits=\"bm25\",\n",
    "                    inputs=[(\"query(q)\", \"tensor<float>(x[384])\")],\n",
    "                    first_phase=\"closeness(field, embedding)\",\n",
    "                    global_phase=GlobalPhaseRanking(\n",
    "                        expression=\"reciprocal_rank_fusion(bm25sum, closeness(field, embedding))\",\n",
    "                        rerank_count=1000,\n",
    "                    ),\n",
    "                ),\n",
    "            ],\n",
    "        )\n",
    "    ],\n",
    "    components=[\n",
    "        Component(\n",
    "            id=\"e5\",\n",
    "            type=\"hugging-face-embedder\",\n",
    "            parameters=[\n",
    "                Parameter(\n",
    "                    \"transformer-model\",\n",
    "                    {\n",
    "                        \"url\": \"https://github.com/vespa-engine/sample-apps/raw/master/simple-semantic-search/model/e5-small-v2-int8.onnx\"\n",
    "                    },\n",
    "                ),\n",
    "                Parameter(\n",
    "                    \"tokenizer-model\",\n",
    "                    {\n",
    "                        \"url\": \"https://raw.githubusercontent.com/vespa-engine/sample-apps/master/simple-semantic-search/model/tokenizer.json\"\n",
    "                    },\n",
    "                ),\n",
    "            ],\n",
    "        )\n",
    "    ],\n",
    ")\n",
    "```\n",
    "\n",
    "Note that the fields `id`, `metadata`, `text`, and `embedding` are required for the integration to work.\n",
    "The schema name must also be `doc`, and the rank profiles must be named `bm25`, `semantic`, and `fusion`.\n",
    "\n",
    "Other than that you are free to modify as you see fit by switching out embedding models, adding more fields, or changing the ranking expressions.\n",
    "\n",
    "For more details, check out this Pyvespa example notebook on [hybrid search](https://pyvespa.readthedocs.io/en/latest/getting-started-pyvespa.html).\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
